import random

import matplotlib.pyplot as plt
import numpy as np


# 代码是从CSDN那个博士的文章里面拷出来的LVoffice
# https://blog.csdn.net/qq_39589163/article/details/136352207?spm=1001.2014.3001.5501
def GS(n, J, P):
    """
    n：机器数
    J：机器矩阵
    P：加工时间矩阵
    """
    MS = []
    time_list = np.zeros(n)
    job_num = len(J)  # 工件数量
    for i in range(job_num):  # i为工件编号
        job_list = J[i]  # job为该工件的工序列表
        for j in range(len(job_list)):  # j为工序编号
            min_time = 99999  # 最小时间初始化
            min_time_index = 0  # 最小位置初始化
            min_time_machine = 0
            machine_list = job_list[j]  # i工件的第j个工序的可选加工机器列表
            process_list = P[i][j]  # i工件的第j个工序的可选加工机器的加工时间列表
            if len(machine_list) == 1:
                MS.append(min_time_index)
                min_time_machine = machine_list[0]
                time_list[min_time_machine] += process_list[min_time_index]
            else:
                for m in range(len(machine_list)):  # m为该工序的可选加工机器列表的第m个机器，机器编号为machine_list[m]
                    temp_time = time_list[machine_list[m]] + process_list[m]
                    if temp_time < min_time:
                        min_time_machine = machine_list[m]
                        min_time = temp_time
                        min_time_index = m
                MS.append(min_time_index)
                time_list[min_time_machine] += process_list[min_time_index]
    return MS


def LS(n, J, P):
    """
    n：机器数
    J：机器矩阵
    P：加工时间矩阵
    """
    MS = []
    job_num = len(J)  # 工件数量
    for i in range(job_num):  # i为工件编号
        time_list = np.zeros(n)
        job_list = J[i]  # job为该工件的工序列表
        for j in range(len(job_list)):  # j为工序编号
            min_time = 99999  # 最小时间初始化
            min_time_index = 0  # 最小位置初始化
            min_time_machine = 0
            machine_list = job_list[j]  # i工件的第j个工序的可选加工机器列表
            process_list = P[i][j]  # i工件的第j个工序的可选加工机器的加工时间列表
            if len(machine_list) == 1:
                MS.append(min_time_index)
                min_time_machine = machine_list[0]
                time_list[min_time_machine] += process_list[min_time_index]
            else:
                for m in range(len(machine_list)):  # m为该工序的可选加工机器列表的第m个机器，机器编号为machine_list[m]
                    temp_time = time_list[machine_list[m]] + process_list[m]
                    if temp_time < min_time:
                        min_time_machine = machine_list[m]
                        min_time = temp_time
                        min_time_index = m
                MS.append(min_time_index)
                time_list[min_time_machine] += process_list[min_time_index]
    return MS


def RS(n, J, P):
    """
    n：机器数
    J：机器矩阵
    P：加工时间矩阵
    """
    MS = []
    job_num = len(J)  # 工件数量
    for i in range(job_num):  # i为工件编号
        job_list = J[i]  # job为该工件的工序列表
        for j in range(len(job_list)):  # j为工序编号
            machine_list = job_list[j]  # i工件的第j个工序的可选加工机器列表
            MS.append(np.random.randint(len(machine_list)))
    return MS


def createInd(n, J, P, split_list=None):
    """
    GLR机器选择方法
    GS：全局选择策略
    LS：局部选择策略
    RS：随机选择策略
    split_list 为三种策略的数量矩阵，如
    split_list = [10, 10, 20]，意为GS、LS各10个个体，RS随机20个个体
    """
    if split_list is None:
        split_list = [10, 10, 20]
    pop = []
    gs = GS(n, J, P)
    ls = LS(n, J, P)
    OS = []
    for i in range(len(J)):
        for _ in range(len(J[i])):
            OS.append(i)
    for _ in range(split_list[0]):
        pop.append(gs + np.random.permutation(OS).tolist())
    for _ in range(split_list[1]):
        pop.append(ls + np.random.permutation(OS).tolist())
    for _ in range(split_list[2]):
        pop.append(RS(n, J, P) + np.random.permutation(OS).tolist())
    return pop


def decode(J, P, s, n):
    """
    function:
    JSP解码函数，用于计算C_max和生成甘特图。
    parameter:
    - J: 机器顺序矩阵
    - P: 加工时间矩阵
    - s: 待解码的序列。
    - n: 机器数
    return:
    - T: 甘特图矩阵。
    - M: 工件排列矩阵。
    - C: 完工时间矩阵。
    """
    job_num = len(J)  # 工件数
    process_num = 0  # 总工序数
    machine_list = []
    max_process_num = 0  # 最大工序数
    # print(J)
    # print(s)
    # machine_list = [item for sublist1 in J for item in sublist1]
    for i in range(job_num):
        machine_list_by_process = s[:len(J[i])] if max_process_num == 0 else s[process_num:process_num + len(J[i])]
        # print(f'i:{i}, list:{machine_list_by_process}, J[i]:{len(J[i])}')
        max_process_num = max(max_process_num, len(J[i]))  # 计算最大工序数
        machine_list.append(machine_list_by_process)
        process_num += len(J[i])  # 计算总工序数
    s = s[process_num:]  # 更新s为后半部分OS
    # print(machine_list, s, process_num, max_process_num)
    # print(f'machine_list:{machine_list}')
    T = [[[0]] for _ in range(n)]
    C = np.zeros((job_num, max_process_num))
    k = np.zeros(job_num, dtype=int)
    # print(s)
    # print(machine_list)
    # print(len(machine_list))
    for job in s:
        # print(f'job:{job}')
        # print(f'm list:{machine_list[job]}')
        machine_index = machine_list[job][k[job]]  # 机器编号
        # print(f'第{job}个工件的第{k[job]}个工序，机器编号是{machine_index}')
        # print(J[job])
        # print(J[job][k[job]])
        machine = J[job][k[job]][machine_index]  # 机器号
        # machine = J[job, k[job]] - 1
        process_time = P[job][k[job]][machine_index]
        # process_time = P[job, k[job]]
        last_job_finish = C[job, k[job] - 1] if k[job] > 0 else 0

        # 寻找机器上的第一个合适空闲时间段
        # print(f'T:{T},machine:{machine}')
        start_time = max(last_job_finish, T[machine][-1][-1])  # 默认在最后一个任务后开始
        insert_index = len(T[machine])  # 默认插入位置在末尾
        for i in range(1, len(T[machine])):
            gap_start = max(T[machine][i - 1][-1], last_job_finish)
            gap_end = T[machine][i][0]
            if gap_end - gap_start >= process_time:
                start_time = gap_start  # 找到合适的起始时间
                insert_index = i  # 更新插入位置
                break
        # print(start_time, process_time)
        end_time = start_time + process_time
        # print(C)
        C[job, k[job]] = end_time
        T[machine].insert(insert_index, [start_time, job, k[job], end_time])
        k[job] += 1

    return T, C


def drawGantt(timeList):
    T = timeList.copy()
    # 创建一个新的图形
    plt.rcParams['font.sans-serif'] = ['SimHei']
    fig, ax = plt.subplots(figsize=(10, 6))

    # 颜色映射字典，为每个工件分配一个唯一的颜色
    color_map = {}
    for machine_schedule in T:
        for task_data in machine_schedule[1:]:
            job_idx, operation_idx = task_data[1], task_data[2]
            if job_idx not in color_map:
                # 为新工件分配一个随机颜色
                color_map[job_idx] = (random.random(), random.random(), random.random())

    # 遍历机器
    for machine_idx, machine_schedule in enumerate(T):
        for task_data in machine_schedule[1:]:
            start_time, job_idx, operation_idx, end_time = task_data
            color = color_map[job_idx]  # 获取工件的颜色

            # 绘制甘特图条形，使用工件的颜色
            ax.barh(machine_idx, end_time - start_time, left=start_time, height=0.4, color=color)

            # 在色块内部标注工件-工序
            label = f'{job_idx}-{operation_idx}'
            ax.text((start_time + end_time) / 2, machine_idx, label, ha='center', va='center', color='white',
                    fontsize=10)

    # 设置Y轴标签为机器名称
    ax.set_yticks(range(len(T)))
    ax.set_yticklabels([f'Machine {i + 1}' for i in range(len(T))])

    # 设置X轴标签
    plt.xlabel("时间")

    # 添加标题
    plt.title("FJSP问题甘特图")

    # 创建图例，显示工件颜色
    legend_handles = []
    for job_idx, color in color_map.items():
        legend_handles.append(plt.Rectangle((0, 0), 1, 1, color=color, label=f'Job {job_idx}'))
    plt.legend(handles=legend_handles, title='工件')

    # # 显示图形，设置分辨率为600
    plt.savefig('FJSP.png', dpi=600)
    # plt.figure(dpi=600) # 调高之后IDE显示不出来了?
    plt.show()


def choice(fitness, k=3, pool=40):
    """
    fitness: 适应度，根据适应度做竞标赛
    k: 每次比较的个数
    pool: 交叉池大小
    """
    n = len(fitness)
    choice_index = []
    for _ in range(pool):
        random_indices = random.sample(range(n), k)
        f_values = [fitness[i] for i in random_indices]
        min_f_value = min(f_values)
        choice_index.append(random_indices[f_values.index(min_f_value)])
    return choice_index


def cross_MS(A, B):
    job_num = len(A)
    random_numbers = random.sample(list(range(job_num)), 2)
    # random_numbers = [0,4]
    rl, rr = min(random_numbers), max(random_numbers)
    afinal = B[:rl] + A[rl:rr + 1] + B[rr + 1:]
    bfinal = A[:rl] + B[rl:rr + 1] + A[rr + 1:]
    return afinal, bfinal


def cross_OS(A, B):
    job_id = list(set(A))
    job_num = len(A)
    # 确保抽取的数量不为0且不等于列表长度
    while True:
        num_to_extract = random.randint(1, len(job_id) - 1)  # 随机抽取的数量，不能为0和等于列表长度
        if num_to_extract > 0 and num_to_extract < len(job_id):
            break
    # 随机抽取元素，S1集合
    S1 = random.sample(job_id, num_to_extract)
    # for list A
    afinal = [None for _ in range(job_num)]
    temp_B = [item for item in B if item not in S1]
    k = 0
    for i in range(job_num):
        if A[i] in S1:
            afinal[i] = A[i]
        else:
            afinal[i] = temp_B[k]
            k += 1
    # for list B
    bfinal = [None for _ in range(job_num)]
    temp_A = [item for item in A if item not in S1]
    k = 0
    for i in range(job_num):
        if B[i] in S1:
            bfinal[i] = B[i]
        else:
            bfinal[i] = temp_A[k]
            k += 1
    return afinal, bfinal


def cross(A, B):
    job_num_all = int(len(A) / 2)
    MS_A = A[:job_num_all]
    OS_A = A[job_num_all:]
    MS_B = B[:job_num_all]
    OS_B = B[job_num_all:]
    MS_A, MS_B = cross_MS(MS_A, MS_B)
    OS_A, OS_B = cross_OS(OS_A, OS_B)
    return MS_A + OS_A, MS_B + OS_B


def mutation_MS(A, P):
    process_num = len(A)
    process_time = [item for sublist1 in P for item in sublist1]  # 拉平P
    r = np.random.randint(process_num)
    random_selection = random.sample(list(range(process_num)), r)
    for i in random_selection:
        A[i] = process_time[i].index(min(process_time[i]))
    return A


def mutation_OS(Ind):
    '''
    函数：变异操作
    内容：使用反转（INV）方法对个体进行变异操作，然后返回变异后的个体。
        具体操作为：随机选择个体中的两个位置，反转这两个位置之间的基因序列。
    '''
    A = Ind.copy()  # 复制个体以避免修改原始个体
    n = len(A)  # 获取工件数
    index1, index2 = random.sample(range(n), 2)  # 随机选择两个不同的索引
    rl, rr = max(index1, index2) + 1, min(index1, index2)  # 确保索引范围正确，避免空序列
    A[rr: rl] = A[rr: rl][::-1]  # 反转选定范围内的基因序列
    return A  # 返回变异后的个体



def mutation(s, P):
    job_num_all = int(len(s) / 2)
    MS = s[:job_num_all]
    OS = s[job_num_all:]
    return mutation_MS(MS, P) + mutation_OS(OS)


def load_data(path):
    # 从文件中读取数据
    with open(path + ".fjs", "r") as file:
        lines = file.readlines()

    # 解析第一行，获取工件数和机器数
    first_line = lines[0].strip().split()
    num_jobs = int(first_line[0])
    num_machines = int(first_line[1])

    # 初始化加工机器矩阵J和加工时间矩阵P
    J = []
    P = []

    # 解析数据行，动态获取工序数并构建J和P矩阵
    for line in lines[1:]:
        # print(line)
        data = line.strip().split()
        if not data:
            continue  # 跳过空行
        # 第一行数据，获取工序数
        num_operations = int(data[0])
        # print(f'工序：{num_operations}')
        data = data[1:]
        job = []
        processing_time = []
        l = 0
        while l < num_operations:
            # print(f'l:{l} data:{data}, num:{num_operations}')
            process_num = int(data[0])
            machine_list = []
            process_list = []
            data = data[1:]
            # print(f'data{data} process_num{process_num}')
            for i in range(process_num):
                machine_list.append(int(data[i * 2]) - 1)  # 机器从0开始编号
                process_list.append(int(data[i * 2 + 1]))
            job.append(machine_list)
            processing_time.append(process_list)
            data = data[i * 2 + 2:]
            l += 1
        J.append(job)
        P.append(processing_time)

    # # 打印结果
    # print("工件数:", num_jobs)
    # print("机器数:", num_machines)
    # print("加工机器矩阵J:", J)
    # print("加工时间矩阵P:", P)
    return J, P, num_machines


if __name__ == "__main__":
    J, P, machine_num = load_data('Mk01')
    # print(J, P, machine_num)
    # J = [[[0, 1, 2, 3, 4], [1, 3]],
    #      [[0, 2, 4], [0, 1, 2], [1, 2, 3, 4]]]
    #
    # P = [[[2, 6, 5, 3, 4], [8, 4]],
    #      [[3, 6, 5], [4, 6, 5], [7, 11, 5, 8]]]
    #
    # machine_num = 5

    pop = createInd(machine_num, J, P, [10, 10, 20])
    pop_size = 40
    # print(pop[0])
    # print(decode(J, P, pop[0], machine_num))
    fitness = [decode(J, P, i, machine_num)[1].max() for i in pop]
    # 初始化最值指标
    best_index = fitness.index(min(fitness))
    best_ind = pop[best_index].copy()
    Cmax = fitness[best_index]
    best_T = decode(J, P, best_ind, machine_num)[0]
    g, gen = 0, 100
    chistory = [Cmax]
    while g < gen:
        g += 1
        # 父代与选择的个体之间进行交叉
        choice_index = choice(fitness, 3, pool=pop_size)
        new_pop = [pop[i] for i in choice_index]
        l = 0
        while l < pop_size / 2:
            if np.random.rand() < 0.7:
                new_pop[l], new_pop[l + 1] = cross(new_pop[l], new_pop[l + 1])
            if np.random.rand() < 0.005:
                new_pop[l] = mutation(new_pop[l], P)
            if np.random.rand() < 0.005:
                new_pop[l + 1] = mutation(new_pop[l + 1], P)
            l += 2
        pop = new_pop
        # 计算种群信息
        fitness = [decode(J, P, i, machine_num)[1].max() for i in pop]

        best_index_temp = fitness.index(min(fitness))
        Cmax_temp = fitness[best_index]
        if Cmax_temp < Cmax:
            Cmax = Cmax_temp
            best_ind = pop[best_index].copy()
            best_T = decode(J, P, best_ind, machine_num)[0]
        chistory.append(Cmax)
    print(f'Cmax:{Cmax}')
    plt.plot(chistory)
    drawGantt(best_T)